package rocksdb

import (
	"crypto/md5"
	"fmt"
	"navidb-node/config"
	"sync"

	"github.com/linxGnu/grocksdb"
)

type ShardManager struct {
	numShards int
	shards    []*grocksdb.DB
	mu        sync.RWMutex
}

func NewShardManager(numShards int, basePath string) (*ShardManager, error) {
	sm := &ShardManager{
		shards:    make([]*grocksdb.DB, numShards),
		numShards: numShards,
	}
	for i := range numShards {
		dbPath := fmt.Sprintf("%s/data_%d", basePath, i)
		db, err := openDB(dbPath)
		if err != nil {
			return nil, fmt.Errorf("failed to open shard %d: %v", i, err)
		}
		sm.shards[i] = db
	}
	return sm, nil
}

func openDB(path string) (*grocksdb.DB, error) {
	bbto := grocksdb.NewDefaultBlockBasedTableOptions()
	bbto.SetBlockCache(grocksdb.NewLRUCache(uint64(config.Conf.Rocks.MemPerShard) * 1024 * 1024)) // 256MB per shard
	opts := grocksdb.NewDefaultOptions()
	opts.SetBlockBasedTableFactory(bbto)
	opts.SetCreateIfMissing(true)
	opts.SetCompression(grocksdb.SnappyCompression)
	opts.SetMaxOpenFiles(5000)
	opts.OptimizeLevelStyleCompaction(uint64(config.Conf.Rocks.TablePerShard) * 1024 * 1024) // 256MB memtable
	return grocksdb.OpenDb(opts, path)
}

func (sm *ShardManager) GetShard(key []byte) (*grocksdb.DB, error) {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	hash := md5.Sum(key)
	shardId := int(hash[0]) % sm.numShards
	return sm.shards[shardId], nil
}

func (sm *ShardManager) Close() {
	sm.mu.Lock()
	defer sm.mu.Unlock()
	for _, db := range sm.shards {
		if db != nil {
			db.Close()
		}
	}
}
